import cv2
import numpy as np
import pyautogui
import pydirectinput
from ultralytics import YOLO
import time


def capture_screen(x, y, width, height):
    """
    捕获指定区域的屏幕截图
    :param x: 区域左上角 x 坐标
    :param y: 区域左上角 y 坐标
    :param width: 区域宽度
    :param height: 区域高度
    :return: 转换为 OpenCV 格式的截图
    """
    screenshot = pyautogui.screenshot(region=(x, y, width, height))
    return cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)


def detect_objects(image, model):
    """
    使用 YOLOv8 模型对图像进行目标检测
    :param image: 待检测的图像
    :param model: 加载好的 YOLO 模型
    :return: 检测结果
    """
    return model(image)


def print_detection_results(results):
    """
    打印检测结果，并返回矿石和自己的位置列表
    :param results: 检测结果
    :return: 矿石位置列表、自己的位置列表和岩浆位置列表
    """
    detection_info = {
        '矿': [],
        '岩浆': [],
        '自己': []
    }
    for r in results:
        boxes = r.boxes
        for box in boxes:
            class_id = r.names[int(box.cls)]
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            if class_id in detection_info:
                detection_info[class_id].append((x1, y1, x2, y2))
    return detection_info['矿'], detection_info['自己'], detection_info['岩浆']


def calculate_distance(box1, box2):
    """
    计算两个边界框中心点的距离
    :param box1: 第一个边界框 (x1, y1, x2, y2)
    :param box2: 第二个边界框 (x1, y1, x2, y2)
    :return: 中心点的距离
    """
    x1_center = (box1[0] + box1[2]) / 2
    y1_center = (box1[1] + box1[3]) / 2
    x2_center = (box2[0] + box2[2]) / 2
    y2_center = (box2[1] + box2[3]) / 2
    return np.sqrt((x1_center - x2_center) ** 2 + (y1_center - y2_center) ** 2)


def is_colliding(box1, box2):
    """
    检查两个边界框是否相交
    :param box1: 第一个边界框 (x1, y1, x2, y2)
    :param box2: 第二个边界框 (x1, y1, x2, y2)
    :return: 如果相交返回 True，否则返回 False
    """
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2
    return not (x2_1 < x1_2 or x2_2 < x1_1 or y2_1 < y1_2 or y2_2 < y1_1)


def avoid_lava(self_position, lava_positions):
    """
    当遇到岩浆时躲开
    :param self_position: 自己的位置
    :param lava_positions: 岩浆位置列表
    """



def move_to_ore(ores, self_positions, lava_positions, x, y, width, height, model):
    """
    控制角色移动到矿石面前
    :param ores: 矿石位置列表
    :param self_positions: 自己的位置列表
    :param lava_positions: 岩浆位置列表
    :param x: 屏幕截图区域左上角 x 坐标
    :param y: 屏幕截图区域左上角 y 坐标
    :param width: 屏幕截图区域宽度
    :param height: 屏幕截图区域高度
    :param model: 加载好的 YOLO 模型
    """
    if ores and self_positions:
        ore = ores[0]
        self_position = self_positions[0]
        distance = calculate_distance(ore, self_position)
        while distance >= 100:
            # 检查是否需要躲开岩浆
            if avoid_lava(self_position, lava_positions):
                # 重新捕获屏幕
                screenshot = capture_screen(x, y, width, height)
                # 重新进行目标检测
                results = detect_objects(screenshot, model)
                # 重新获取矿石、自己和岩浆的位置
                ores, self_positions, lava_positions = print_detection_results(results)
                if not ores or not self_positions:
                    print("未检测到矿石或自己的位置，停止移动")
                    break
                ore = ores[0]
                self_position = self_positions[0]
                # 重新计算距离
                distance = calculate_distance(ore, self_position)
                continue

            ore_center_x = (ore[0] + ore[2]) / 2
            self_center_x = (self_position[0] + self_position[2]) / 2
            ore_center_y = (ore[1] + ore[3]) / 2
            self_center_y = (self_position[1] + self_position[3]) / 2

            # 临时存储可能的移动后的位置
            new_self_position = list(self_position)

            if ore_center_x > self_center_x:
                new_self_position[0] += 80  # 假设每次移动 100 像素
                new_self_position[2] += 80
                can_move = True
                for lava in lava_positions:
                    if is_colliding(new_self_position, lava):
                        can_move = False
                        break
                if can_move:
                    pydirectinput.keyDown('right')
                    time.sleep(0.1)
                    pydirectinput.keyUp('right')
                new_self_position = list(self_position)  # 恢复原始位置

            elif ore_center_x < self_center_x:
                new_self_position[0] -= 80
                new_self_position[2] -= 80
                can_move = True
                for lava in lava_positions:
                    if is_colliding(new_self_position, lava):
                        can_move = False
                        break
                if can_move:
                    pydirectinput.keyDown('left')
                    time.sleep(0.1)
                    pydirectinput.keyUp('left')
                new_self_position = list(self_position)

            if ore_center_y > self_center_y:
                new_self_position[1] += 80
                new_self_position[3] += 80
                can_move = True
                for lava in lava_positions:
                    if is_colliding(new_self_position, lava):
                        can_move = False
                        break
                if can_move:
                    pydirectinput.keyDown('down')
                    time.sleep(0.1)
                    pydirectinput.keyUp('down')
                new_self_position = list(self_position)

            elif ore_center_y < self_center_y:
                new_self_position[1] -= 10
                new_self_position[3] -= 10
                can_move = True
                for lava in lava_positions:
                    if is_colliding(new_self_position, lava):
                        can_move = False
                        break
                if can_move:
                    pydirectinput.keyDown('up')
                    time.sleep(0.1)
                    pydirectinput.keyUp('up')
                new_self_position = list(self_position)

            # 重新捕获屏幕
            screenshot = capture_screen(x, y, width, height)
            # 重新进行目标检测
            results = detect_objects(screenshot, model)
            # 重新获取矿石、自己和岩浆的位置
            ores, self_positions, lava_positions = print_detection_results(results)
            if ores and self_positions:
                ore = ores[0]
                self_position = self_positions[0]
                # 重新计算距离
                distance = calculate_distance(ore, self_position)
            else:
                print("未检测到矿石或自己的位置，停止移动")
                break

        if distance < 100:
            # 再次确认距离
            ore_center_x = int((ore[0] + ore[2]) / 2)
            ore_center_y = int((ore[1] + ore[3]) / 2 + 80)

            new_distance = calculate_distance(ore, self_position)
            if new_distance < 100:
                print(f"到达矿石位置，点击挖矿，矿石坐标: ({ore_center_x}, {ore_center_y})")
                pydirectinput.click(ore_center_x, ore_center_y)

            else:
                print("距离矿石仍较远，不进行点击操作")


def main():
    global x, y, width, height, model
    x = 0
    y = 87
    width = 900
    height = 570

    # 加载模型
    model = YOLO('best.pt')

    while True:
        # 捕获屏幕截图
        screenshot = capture_screen(x, y, width, height)
        # yolo 检测
        results = detect_objects(screenshot, model)
        # 获取并打印自己、矿石和岩浆的位置
        ores, self_positions, lava_positions = print_detection_results(results)

        # 移动到矿石面前
        move_to_ore(ores, self_positions, lava_positions, x, y, width, height, model)

        for r in results:
            im_array = r.plot()
            cv2.imshow('YOLOv8 Inference', im_array)

        # 可根据需要添加退出条件
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
